/*
 * Copyright 2008-2010 Sun Microsystems, Inc.  All Rights Reserved.
 * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
 *
 * This code is free software; you can redistribute it and/or modify it
 * under the terms of the GNU General Public License version 2 only, as
 * published by the Free Software Foundation.  Sun designates this
 * particular file as subject to the "Classpath" exception as provided
 * by Sun in the LICENSE file that accompanied this code.
 *
 * This code is distributed in the hope that it will be useful, but WITHOUT
 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
 * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
 * version 2 for more details (a copy is included in the LICENSE file that
 * accompanied this code).
 *
 * You should have received a copy of the GNU General Public License version
 * 2 along with this work; if not, write to the Free Software Foundation,
 * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
 *
 * Please contact Sun Microsystems, Inc., 4150 Network Circle, Santa Clara,
 * CA 95054 USA or visit www.sun.com if you need additional information or
 * have any questions.
 */

package com.sun.btrace.agent;

import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.ObjectInputStream;
import java.io.ObjectOutputStream;
import java.io.OutputStream;
import java.lang.instrument.ClassDefinition;
import java.lang.instrument.Instrumentation;
import java.net.Socket;
import java.util.Arrays;
import java.util.Enumeration;
import java.util.Iterator;
import java.util.Map;
import java.util.Set;
import java.util.TreeMap;
import java.util.jar.JarEntry;
import java.util.jar.JarFile;
import java.util.jar.JarOutputStream;

import com.sun.btrace.BTraceRuntime;
import com.sun.btrace.BTraceUtils;
import com.sun.btrace.comm.ClassPairBody;
import com.sun.btrace.comm.Command;
import com.sun.btrace.comm.ConfBody;
import com.sun.btrace.comm.ConfCommand;
import com.sun.btrace.comm.ErrorCommand;
import com.sun.btrace.comm.EventCommand;
import com.sun.btrace.comm.ExitCommand;
import com.sun.btrace.comm.InstrumentCommand;
import com.sun.btrace.comm.OkayCommand;
import com.sun.btrace.comm.WireIO;
import com.sun.btrace.comm.ConfCommand.OPERATOR;

/**
 * Represents a remote client communicated by socket.
 * 
 * @author A. Sundararajan
 */
class RemoteClient extends Client {
	private volatile Socket sock;
	private volatile ObjectInputStream ois;
	private volatile ObjectOutputStream oos;

	RemoteClient(Instrumentation inst, Socket sock) throws IOException {
		super(inst);
		this.sock = sock;
		this.ois = new ObjectInputStream(sock.getInputStream());
		this.oos = new ObjectOutputStream(sock.getOutputStream());
	}

	public void startServer() throws IOException {

		Command cmd = WireIO.read(ois);
		if (cmd.getType() == Command.INSTRUMENT) {
			if (debug)
				Main.debugPrint("got instrument command");
			Class btraceClazz = loadClass((InstrumentCommand) cmd);
			if (btraceClazz == null) {
				throw new RuntimeException("can not load BTrace class");
			}
		} else {
			errorExit(new IllegalArgumentException(
					"expecting instrument command!"));
			throw new IOException("expecting instrument command!");
		}
		Thread cmdHandler = new Thread(new Runnable() {
			public void run() {
				BTraceRuntime.enter();
				while (true) {
					try {
						Command cmd = WireIO.read(ois);
						switch (cmd.getType()) {
						case Command.EXIT: {
							ExitCommand ecmd = (ExitCommand) cmd;
							if (debug)
								Main.debugPrint("received exit command");
							BTraceRuntime.leave();
							BTraceRuntime.enter(getRuntime());
							try {
								if (debug)
									Main.debugPrint("calling BTraceUtils.exit()");
								BTraceUtils.Sys.exit(ecmd.getExitCode());
							} catch (Throwable th) {
								if (debug)
									Main.debugPrint(th);
								BTraceRuntime.handleException(th);
							}
							finally{
								closeAll();
							}
							return;
						}
						case Command.EVENT: {
							getRuntime().handleEvent((EventCommand) cmd);
							break;
						}
						default:
							if (debug)
								Main.debugPrint("received " + cmd);
							// ignore any other command
						}
					} catch (Exception exp) {
						if (debug)
							Main.debugPrint(exp);
						break;
					}
				}
			}
		});
		cmdHandler.setDaemon(true);
		if (debug)
			Main.debugPrint("starting client command handler thread");
		cmdHandler.start();
	}

	public boolean beforeStartServer() throws IOException {
		boolean isRight = false;
		String ip = sock.getLocalAddress().getHostAddress();
		Main.debugPrint("client connect from " + ip);
		// if isRight == false, this socket will stop.

		Command cmd = WireIO.read(ois);
		Main.debugPrint("accept client, get cmd " + cmd.getType());
		if( cmd.getType() != Command.CONF){
			WireIO.write(oos, new ErrorCommand(new Exception("need configure command, not "+ cmd.getType()))) ;
			return false ;
		}
		ConfCommand conf = (ConfCommand)cmd ;
		ConfBody body = conf.getBody() ;
		ClassPairBody[] cbody = conf.getCbody() ;

		if( ! checkAuth(body)&& conf.getOperator()!= OPERATOR.SETPASSWORD && conf.getOperator()!= OPERATOR.RESETPASSWORD){
			WireIO.write(oos, new ErrorCommand(new Exception("Wrong password!!!")));
			return false;
		}
		switch(conf.getOperator()){
		case SETPASSWORD:
			Main.setPassword(body.getPassword());
			Main.debugPrint("client set password to ***"+body.getPassword());
			WireIO.write(oos, new OkayCommand());
			break;
		case RESETPASSWORD:
				Main.debugPrint("reset pasword");
				String host=sock.getInetAddress().getHostAddress();
				boolean local=host.equals("localhost") || host.equals("127.0.0.1");
				if(local){
					Main.setPassword(body.getNewPassword());
					WireIO.write(oos, new OkayCommand());
				}else{
					if( checkAuth(body)){
						Main.setPassword(body.getNewPassword());
						WireIO.write(oos, new OkayCommand());
					}
					else{
						WireIO.write(oos, new ErrorCommand(new Exception ("password is not right")));
					}
				}
				
				break;
		case GETCONF:
				ConfBody rbody = new ConfBody();
				getConfigureFromAgent(rbody);
				ConfCommand r = new ConfCommand();
				r.setOperator(OPERATOR.GETCONF);
				r.setBody(body);
				WireIO.write(oos, r);
				break;
		case SETCONF:
				setConfigureForAgent(conf.getBody());
				WireIO.write(oos, new OkayCommand());
				break;
		case LOGIN:
			WireIO.write(oos, new OkayCommand());

			return true ;
		case REFINECLASS:
			if( ! this.inst.isRedefineClassesSupported()){
				WireIO.write(oos, new ErrorCommand(new Exception("refine class is not support!")) );
				break ;
			}
			try {
				ClassDefinition[]cls = new ClassDefinition[cbody.length] ;
				for( int i=0; i<cls.length; ++i){
					cls[i] = new ClassDefinition(Class.forName(cbody[i].getFullName()), cbody[i].getCode());
				}
				this.inst.redefineClasses(cls);
				WireIO.write(oos, new OkayCommand());
				Main.debugPrint("OK to refine the class: "+ Arrays.toString(cbody) );
			} catch (Exception e) {
				WireIO.write(oos, new ErrorCommand(e)) ;
				Main.debugPrint("fail to refine the class: "+  Arrays.toString(cbody)) ;
			}
			
			break;
		case REFINECLASS2FILE:
			if( ! this.inst.isRedefineClassesSupported()){
				WireIO.write(oos, new ErrorCommand(new Exception("refine class is not support!")) );
				break ;
			}
			try {
				ClassDefinition[]cls = new ClassDefinition[cbody.length] ;
				for( int i=0; i<cls.length; ++i){
					cls[i] = new ClassDefinition(Class.forName(cbody[i].getFullName()), cbody[i].getCode());
				}
				this.inst.redefineClasses(cls);
				Main.debugPrint("refine the classes in memory: " +   Arrays.toString(cbody)) ;
				Main.debugPrint("begin to save class to file." );

				for( int i=0; i<cls.length; ++i){
					String file= Thread.currentThread().getContextClassLoader().getResource(cbody[i].getFullName()+".class").toString();
					if( file.startsWith("jar:")){
						String jar= file.substring("jar:file:/".length(), file.lastIndexOf("!")) ;
						String cfile=cbody[i].getFullName().replace(".", "/")+".class" ;
						editJar(jar, cfile, cbody[i].getCode());
					}
					else{
						FileOutputStream out = new FileOutputStream(file+cbody[i].getFullName().replace(".", "/")+".class") ;
						out.write(cbody[i].getCode());
					}
					
				}	
				WireIO.write(oos, new OkayCommand());
				Main.debugPrint("OK to refine the class: "+  Arrays.toString(cbody)) ;
				
			} catch (Exception e) {
				WireIO.write(oos, new ErrorCommand(e)) ;
				Main.debugPrint("fail to refine the class: "+  Arrays.toString(cbody)) ;
			}
			
			break;
			
		case ATTACHJAR:
			try{
				for( int i=0; i<cbody.length; ++i){
					String file= cbody[i].getFullName();
					OutputStream os = new FileOutputStream(new File(file));
					os.write(cbody[i].getCode());
					JarFile jarfile = new JarFile(file);
					this.inst.appendToSystemClassLoaderSearch(jarfile) ;
				}
				WireIO.write(oos, new OkayCommand());
				Main.debugPrint("OK to append jars to jvm: "+  Arrays.toString(cbody)) ;
			}
			catch (Exception e) {
				WireIO.write(oos, new ErrorCommand(e)) ;
				Main.debugPrint("fail to add the class: "+  Arrays.toString(cbody)) ;
			}
			
			break;
		default:
			break;
		}
 
		return isRight;
	}
	public  void editJar(String file,String entryName,byte[] data){ 
	 try {
		   JarFile jf = new JarFile(file);
		   TreeMap tm = new TreeMap();
		   Enumeration es = jf.entries();
		   while (es.hasMoreElements()) {
		    JarEntry je = (JarEntry) es.nextElement();
		    byte[] b = getByte(jf.getInputStream(je));
		    tm.put(je.getName(), b);
		   }
		   JarOutputStream out = new JarOutputStream(
		     new FileOutputStream(file));
		   Set set = tm.entrySet();
		   Iterator it = set.iterator();
		   boolean has = false;
		   while (it.hasNext()) {
		    Map.Entry me = (Map.Entry) it.next();
		    String name = (String) me.getKey();
		    JarEntry jeNew = new JarEntry(name);
		    out.putNextEntry(jeNew);
		    byte[] b;
		    if (name.equals(entryName)) {
		     b = data;
		     has = true;
		    } else
		     b = (byte[]) me.getValue();
		    out.write(b, 0, b.length);
		   }
		   if (!has) {
		    JarEntry jeNew = new JarEntry(entryName);
		    out.putNextEntry(jeNew);
		    out.write(data, 0, data.length);
		   }
		   out.finish();
		   out.close();
		   jf.close();
		  } catch (Exception e) {
		   e.printStackTrace();
		  }
		 }
	public   byte[] getByte(java.io.InputStream s) {
		  byte[] buffer = new byte[0];
		  byte[] chunk = new byte[4096];
		  int count;
		  try {
		   while ((count = s.read(chunk)) >= 0) {
		    byte[] t = new byte[buffer.length + count];
		    System.arraycopy(buffer, 0, t, 0, buffer.length);
		    System.arraycopy(chunk, 0, t, buffer.length, count);
		    buffer = t;
		   }
		   s.close();
		  } catch (Exception e) {
		  }
		  return buffer;
		 }
	protected boolean checkAuth(ConfBody conf) {
		/*
		 * below part is use ip to check user's right;
		 * now it will use psw to check
		 */
		/*
		String ip = sock.getLocalAddress().getHostAddress();
		Main.debugPrint("client connect from " + ip);
		if (ip.equals("localhost") || ip.equals("127.0.0.1")) {
			return true;
		}
		return false;
		*/
		Main.debugPrint(Main.getPassword()+"?="+conf.getPassword());
		//check the psw at current;
		return Main.getPassword().equals(conf.getPassword());
	}

	public void setConfigureForAgent(ConfBody body) {
		Main.debugPrint("set configure for the agent.");
		Main.setDebugMode(body.isDebugMode());
		Main.setTrackRetransforms(body.isTrackRetransforms());
		Main.setUnsafeMode(body.isUnsafeMode());
		Main.setDumpClasses(body.isDumpClasses());
		Main.setDumpDir(body.getDumpDir());
		Main.setScriptOutputFile(body.getScriptOutputFile());
//		Main.setPassword(body.getPassword());
	}

	public void getConfigureFromAgent(ConfBody body) {
		Main.debugPrint("get configure for the agent.");

		body.setDebugMode(Main.isDebugMode());
		body.setTrackRetransforms(Main.isTrackRetransforms());
		body.setUnsafeMode(Main.isUnsafeMode());
		body.setDumpClasses(Main.isDumpClasses());
		body.setDumpDir(Main.getDumpDir());
		body.setScriptOutputFile(Main.getScriptOutputFile());
//		body.setPassword(Main.getPassword());
	}

	public void onCommand(Command cmd) throws IOException {
		if (oos == null) {
			throw new IOException("no output stream");
		}
		oos.reset();
		switch (cmd.getType()) {
		case Command.EXIT:
			if (debug)
				Main.debugPrint("client " + getClassName() + ": got " + cmd);
			WireIO.write(oos, cmd);
			onExit(((ExitCommand) cmd).getExitCode());
			break;
		default:
			if (debug)
				Main.debugPrint("client " + getClassName() + ": got " + cmd);
			WireIO.write(oos, cmd);
		}
	}

	protected synchronized void closeAll() throws IOException {
		if (oos != null) {
			oos.close();
		}
		if (ois != null) {
			ois.close();
		}
		if (sock != null) {
			sock.close();
		}
		this.runtime = null ;
		Main.debugPrint("close the socket resource for the client") ;
	}
}